# coding = utf-8

'''
转换数据到chengkun格式的图像
'''

import os
import SimpleITK as sitk
from pathlib2 import Path
import numpy as np

def transfer_data():
    list_root_path = "F:\Dataset\Liver\LITS"
    save_root_path = "F:\Dataset\Liver\chengkun_only_liver"

    for i in range(131):
        raw_path = os.path.join(list_root_path, "volume-{}.nii".format(i))
        mask_path = os.path.join(list_root_path, "segmentation-{}.nii".format(i))

        raw_data = sitk.GetArrayFromImage(sitk.ReadImage((raw_path)))
        mask_data = sitk.GetArrayFromImage(sitk.ReadImage(mask_path))

        raw_data[raw_data >= 250] = 250
        raw_data[raw_data <= -200] = -200
        raw_data = (raw_data + 200) / 450

        mask_data = mask_data.astype(np.uint8)

        case_id = "case_{}".format(str(i).zfill(5))

        destination_patient_path = os.path.join(save_root_path, case_id)
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        index = 0
        for j in range(raw_data.shape[0]):
            if np.max(mask_data[j]) == 0:
                continue
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(index).zfill(3))), raw_data[j])
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(index).zfill(3))), mask_data[j])
            index += 1





if __name__ == '__main__':
    transfer_data()